import torch
from PIL import Image
from torchvision import transforms
from train import VGG16net

import time

from utils import getClassify

name2label = getClassify(177)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
resize = 224
tf = transforms.Compose([  # 常用的数据变换器
    lambda x: Image.open(x).convert('RGB'),  # string path= > image data
    # 这里开始读取了数据的内容了
    transforms.Resize(  # 数据预处理部分
        (int(resize * 1.25), int(resize * 1.25))),
    transforms.RandomRotation(15),
    transforms.CenterCrop(resize),  # 防止旋转后边界出现黑框部分
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])


def prediect(img_path):
    model = VGG16net().to(device)
    model.load_state_dict(torch.load('model/model.pth'))
    # net = net.to(device)
    with torch.no_grad():
        img = tf(img_path).unsqueeze(0)
        img_ = img.to(device)
        start = time.time()
        outputs = model(img_)
        _, predicted = torch.max(outputs, 1)
        predicted_number = predicted[0].item()
        end = time.time()
        print('this picture maybe :', list(name2label.keys())[list(name2label.values()).index(predicted_number)])
        print('FPS:', 1 / (end - start))
    return list(name2label.keys())[list(name2label.values()).index(predicted_number)]



if __name__ == '__main__':
    # acc = []
    # for j in range(0,3):
    #     n = 0.0
    #     for i in range(1,181):
    #         a=prediect('./train/img/'+str(j)+'/'+str(i)+'.jpg')
    #         print(a)
    #         if int(a)==j:
    #             print(1)
    #             n=n+1
    #     print(n/180)
    #     acc.append(n/180)
    #     print(n)
    # print(acc)
    a = prediect('./test/1.jpg')
    b = prediect('./test/2.jpg')
    c = prediect('./test/img_1.png')
    d = prediect('./test/img.png')
